import os
import operator
import asyncio
from dotenv import load_dotenv
from typing import Annotated, Any, TypedDict

from langchain_core.messages import AnyMessage, HumanMessage, AIMessage
from langchain_openai import ChatOpenAI

from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import MemorySaver

# --- 1. 环境准备 ---
load_dotenv()

# --- 2. 状态定义 ---
class AgentState(TypedDict):
    messages: Annotated[list[AnyMessage], operator.add]

# --- 3. 节点定义 ---
# a. 核心变化：一个能够"异步流式产出"的 Agent 节点
llm = ChatOpenAI(model="qwen-plus-latest", base_url=os.getenv("OPENAI_BASE_URL"))

# 2. 定义 Agent 节点
# 这个节点现在是一个标准的异步函数，而不是生成器。
# LangGraph 会自动处理来自内部 LLM 调用的流式令牌。
async def streaming_agent_node(state: AgentState):
    """
    一个标准的异步节点。它调用 LLM，而 LangGraph 在外部处理流式传输。
    """
    print("\n---AGENT: 思考中...---")
    
    # 直接调用 .ainvoke() 来获取完整响应
    # LangGraph 将自动捕获此调用产生的令牌流
    response = await llm.ainvoke(state["messages"])
    
    # 返回完整的消息以更新状态
    return {"messages": [response]}

# --- 4. 构建图 ---
# 图的结构非常简单，只有一个节点
workflow = StateGraph(AgentState)
workflow.add_node("agent", streaming_agent_node)
workflow.set_entry_point("agent")
workflow.set_finish_point("agent")

# --- 5. 编译图 ---
# 流式输出不强制要求检查点，但如果需要记忆和持久化，加上也无妨
checkpointer = MemorySaver()
app = workflow.compile(checkpointer=checkpointer)

# --- 6. 消费流式输出 ---
async def main():
    """主异步函数，用于运行和消费流式输出"""
    question = "写一个关于一只生活在赛博朋克城市的猫的短篇故事。"
    
    config = {"configurable": {"thread_id": "user_session_4"}}
    inputs = {"messages": [HumanMessage(content=question)]}
    
    print("---开始接收流式输出---")
    
    # astream_events 是专门用于流式获取事件的API
    # 我们只关心 'on_chat_model_stream' 事件，它包含了模型生成的 token
    async for event in app.astream_events(inputs, config=config, version="v2"):
        kind = event["event"]
        if kind == "on_chat_model_stream":
            content = event["data"]["chunk"].content
            if content:
                # 通过 end="" 和 flush=True 实现打字机效果
                print(content, end="", flush=True)

    print("\n\n---流式输出结束---")

if __name__ == "__main__":
    # 使用 asyncio.run() 来启动我们的主异步函数
    asyncio.run(main()) 